Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
三阶线性注意力:三维循环状态用于长上下文序列建模
🧭 线性注意力 · 长上下文状态扩容 ★ 深度解读
Massachusetts Institute of Technology · MIT-IBM Computing Research Lab
将线性注意力的矩阵状态扩展为三阶张量,用第二个 key 维度以极少参数成倍扩大状态容量,显著提升长上下文语言建模与 recall
📌 工作参考价值(4 点)
- 电商/推荐用户长期行为序列建模:用户行为序列长且稀疏,Transformer KV cache 显存贵;可把 user encoder 或长期兴趣提取器替换为 Triadic GDN/sGLA,用 E=4~8 扩大状态,参数只增加约 1%,线上保持 O(1) 状态,对历史兴趣 recall 提升明显。second key/query 建议用 softplus/sigmoid 非负激活,避免切片间抵消。
- Agent 长期记忆/长上下文工具结果:如果 Agent 用线性注意力做记忆,可在长上下文扩展阶段把 E=1 upcycle 到 E=8,不必从头训练;论文显示 upcycling 可恢复约 50%-75% 从 scratch 的收益,适合已有模型低成本升级。
- 混合架构取舍:搜推广 rerank 或 Agent 模型若采用 linear+softmax 混合,优先扩大 linear state 而不是 softmax GQA 的 KV cache;论文在 3:1 混合中,triadic E=4 比 GQA-4 在长上下文 PG19 和 NIAH 更好,且 64k 显存少约一半。
- 工程实现:chunkwise-parallel 实现中,把三阶 state 沿 value 轴切成 32 列块放到 registers,避免 SM 持有完整 state;joint key 内积分利用 Kronecker 结构降维,E 维只增加少量 C^2E 计算,长上下文效率优于 Transformer。